#pragma once

#include "echo.hpp"
#include "limbo/errc.h"
#include "limbo/ip/ipv4_address.h"
#include "state.hpp"
#include <alloca.h>
#include <cstdint>

namespace limbo::icmp {

template <typename IcmpState> struct DestinationUnreacheable {
  using ContextState = DestinationUnreacheable;
  using Context = DestinationUnreacheable;
  using LowerState = IcmpState;
  using LowerContext = typename IcmpState::Context;
  using LowerPacket = typename LowerState::Packet;
  using IpState = typename IcmpState::LowerState;
  using IpContext = typename IpState::Context;
  using IpPacket = typename IpState::Packet;
  using LinkContext = typename IpContext::LowerContext;
  using Customization = customization<details::base_state_t<IcmpState>>;

  static constexpr uint8_t SZ = 4;

  inline void init(const ip::IPv4Address &source, uint8_t tos = 0,
                   uint8_t ttl = 64) {
    if constexpr (details::has_lower_state_v<IpState>) {
      ip_ctx.link_context = &link_ctx_wrapper.ctx;
    }
    ip_ctx.source = source;
    ip_ctx.proto = ip::Proto::icmp;
    ip_ctx.ttl = ttl;
    ip_ctx.tos = tos;
    lower_state = nullptr;
  }

  Result recv(Chunk buff, LowerState *icmp_state) {
    lower_state = icmp_state;
    return Result::make_consumed(this, buff);
  }

  LowerPacket &get_parsed() { return lower_state->get_parsed(); }

  inline bool select(const LowerPacket *packet) {
    return (packet->type == Type::destination_unreachable) &&
           (packet->container->destination == ip_ctx.source);
  }

  template <typename Stack>
  Result send(Stack &stack, Chunk buff, const IpPacket &ip_packet, uint8_t code,
              uint16_t mtu = {}) {

    auto payload_sz = std::min(16u, ip_packet.payload.size());
    auto sz = ip_packet.header_length * sizeof(uint32_t) + payload_sz;
    if (buff.size() < sz + 4) {
      auto code = (uint32_t)Errc::no_enough_space;
      return Result::make_error(code, this);
    }

    ip_ctx.destination = ip_packet.source;
    icmp_ctx.code = code;

    uint8_t *payload_raw = static_cast<uint8_t *>(alloca(sz));
    auto ptr = payload_raw;
    *ptr++ = 0;
    *ptr++ = 0; /* unused */
    auto mtu_be = order::native_to_big(mtu);
    Customization::copy(ptr, &mtu_be, sizeof(mtu_be));
    auto payload = Chunk(payload_raw, 4);
    auto ip_chunk = Chunk(ip_packet.packet_start, sz);
    Customization::copy(payload, ip_chunk);
    return stack.send(buff, *this, payload);
  }

  details::LinkContextWrapper<LinkContext> link_ctx_wrapper;
  IpContext ip_ctx;
  LowerContext icmp_ctx = {Type::destination_unreachable, 0, &ip_ctx};
  LowerState *lower_state;
};

template <typename State>
inline auto *get_lower_context(DestinationUnreacheable<State> &ctx) {
  return &ctx.icmp_ctx;
}

template <typename State>
inline auto &get_link_context(DestinationUnreacheable<State> &ctx) {
  using IpState = typename State::LowerState;
  static_assert(
      details::has_lower_state_v<IpState>,
      "the function has meaning only when there is link state in the stack");
  auto icmp_ctx = get_lower_context(ctx);
  auto ip_ctx = get_lower_context(*icmp_ctx);
  auto link_ctx = get_lower_context(*ip_ctx);
  return *link_ctx;
}

} // namespace limbo::icmp
